import Driver
import FrontEnd
import IR
import Utils
import XCTest

extension Diagnostic {

  /// A test annotation that announces `self` should be expected.
  fileprivate var expectation: TestAnnotation {
    TestAnnotation(
      in: site.file.url,
      atLine: site.start.line.number,
      parsing: "diagnostic " + message
    )
  }

}

extension Result {

  /// The success value, if any, and `nil` otherwise.
  var success: Success? {
    if case .success(let r) = self { return r } else { return nil }
  }

  /// The failure value, if any, and `nil` otherwise.
  var failure: Failure? {
    if case .failure(let r) = self { return r } else { return nil }
  }

}

/// The expected outcome of a test.
public enum ExpectedTestOutcome: Sendable {

  case success, failure

}

extension XCTContextHylo {

  // Same as `XCTContext.runActivity(named:block:)`, but without the `@MainActor` requirement.
  public static func runActivityNonIsolated<Result>(
    named debugName: String,
    block: (XCTActivityHylo) throws -> Result
  ) rethrows -> Result {
    try block(XCTActivityHylo(debugName: debugName))
  }

}
extension XCTestCase {

  /// The effects of running the `processAndCheck` parameter to `checkAnnotatedHyloFiles`.
  private struct ProcessingEffects: @unchecked Sendable {

    /// Test failures generated by processing.
    let testFailures: [XCTIssue]

    /// Hylo diagnostics generated by processing.
    let diagnostics: DiagnosticSet

  }

  /// Applies `processAndCheck` to `hyloToTest` and the subset of its annotations whose commands
  /// match `checkedCommands`, recording resulting XCTest failures along with any additional
  /// failures where the effects of processing don't match the its annotation commands ("//!
  /// ... diagnostic ..."), and returning any error thrown by `processAndCheck`.
  ///
  /// - Parameters:
  ///   - checkedCommands: the annotation commands to be validated by `processAndCheck`.
  ///   - processAndCheck: applies some compilation phases to `file`, updating `diagnostics`
  ///     with any generated diagnostics, then checks `annotationsToCheck` against the results,
  ///     returning corresponding test failures. Throws an `Error` if any phases failed.
  private func checkAnnotations(
    in hyloToTest: SourceFile,
    checkingAnnotationCommands checkedCommands: Set<String> = [],
    _ processAndCheck: (
      _ file: SourceFile,
      _ annotationsToCheck: ArraySlice<TestAnnotation>,
      _ diagnostics: inout DiagnosticSet
    ) throws -> [XCTIssue]
  ) -> Error? {
    var annotations = TestAnnotation.parseAll(from: hyloToTest)

    // Separate the annotations to be checked by default diagnostic annotation checking from
    // those to be checked by `processAndCheck`.
    let p = annotations.partition(by: { checkedCommands.contains($0.command) })
    let (diagnosticAnnotations, processingAnnotations) = (annotations[..<p], annotations[p...])

    var diagnostics = DiagnosticSet()
    var thrownError: Error? = nil

    let failures = XCTContextHylo.runActivityNonIsolated(
      named: hyloToTest.baseName,
      block: { activity in
        let r = Result { try processAndCheck(hyloToTest, processingAnnotations, &diagnostics) }
        thrownError = r.failure

        return failuresToReport(
          effectsOfProcessing: .init(
            testFailures: r.success ?? [],
            diagnostics: diagnostics),
          unhandledAnnotations: diagnosticAnnotations)
      })

    for f in failures {
      record(f)
    }

    return thrownError
  }

  /// Applies `process` to the ".hylo" file at the given path and reports XCTest failures where the
  /// effects of processing don't match the file's annotation commands ("//! ... diagnostic ...").
  ///
  /// - Parameters:
  ///   - process: applies some processing to `file`, updating `diagnostics` with any generated
  ///     diagnostics. Throws an `Error` if processing failed.
  ///   - expectSuccess: true if an error from `process` represents a test failure, false if the
  ///     lack of an error represents a test failure; nil if that information is to be derived
  ///     from the contents of the file.
  private func checkAnnotatedHyloFileDiagnostics(
    inFileAt hyloFilePath: String,
    expecting expectation: ExpectedTestOutcome,
    _ process: (_ file: SourceFile, _ diagnostics: inout DiagnosticSet) throws -> Void
  ) throws {
    let f = try SourceFile(at: hyloFilePath)

    // FIXME: clarify/explain this code
    let thrownError = checkAnnotations(in: f, checkingAnnotationCommands: []) {
      (f, annotationsToHandle, diagnostics) in
      assert(annotationsToHandle.isEmpty)
      try process(f, &diagnostics)
      return []
    }

    if (thrownError == nil) != (expectation == .success) {
      record(XCTIssue(unexpectedOutcomeDiagnostic(thrownError: thrownError, at: f.wholeRange)))
    }
  }

  /// Returns the diagnostic of an unexpected outcome with given `thrownError` reported at `s`.
  private func unexpectedOutcomeDiagnostic(thrownError: Error?, at s: SourceRange) -> Diagnostic {
    if let e = thrownError {
      return .error("success was expected, but processing failed with thrown error: \(e)", at: s)
    } else {
      return .error("processing succeeded, but failure was expected", at: s)
    }
  }

  /// Given the effects of processing, the annotations not specifically handled by `processAndCheck`
  /// above, returns the final set of test failures to be reported to XCTest.
  private func failuresToReport(
    effectsOfProcessing processing: ProcessingEffects,
    unhandledAnnotations: ArraySlice<TestAnnotation>
  ) -> [XCTIssue] {
    var testFailures = processing.testFailures

    var diagnosticsByExpectation = Dictionary(
      grouping: processing.diagnostics.elements, by: \.expectation)

    func fail(_ expectation: TestAnnotation, _ message: String) {
      testFailures.append(expectation.failure(message))
    }

    for a in unhandledAnnotations {
      switch a.command {
      case "diagnostic":
        if diagnosticsByExpectation[a]?.popLast() != nil {
        } else {
          fail(a, "missing expected diagnostic\(a.argument.map({": '\($0)'"}) ?? "")")
        }
      case "expect-failure": do {}
      case "expect-success": do {}
      default:
        fail(a, "unexpected test command: '\(a.command)'")
      }
    }

    testFailures += diagnosticsByExpectation.values.joined().lazy.map {
      XCTIssue(.error("unexpected diagnostic: '\($0.message)'", at: $0.site, notes: $0.notes))
    }
    return testFailures
  }

  /// Calls `compileAndRun` with optimizations disabled.
  @nonobjc
  public func compileAndRun(
    _ hyloFilePath: String, extending p: TypedProgram, expecting expectation: ExpectedTestOutcome
  ) throws {
    try compileAndRun(hyloFilePath, withOptimizations: false, extending: p, expecting: expectation)
  }

  /// Calls `compileAndRun` with optimizations enabled.
  @nonobjc
  public func compileAndRunOptimized(
    _ hyloFilePath: String, extending p: TypedProgram, expecting expectation: ExpectedTestOutcome
  ) throws {
    try compileAndRun(hyloFilePath, withOptimizations: true, extending: p, expecting: expectation)
  }

  /// Compiles and runs the hylo file at `hyloFilePath`, applying program optimizations iff
  /// `withOptimizations` is `true`, and `XCTAssert`ing that diagnostics and exit codes match
  /// annotated expectations.
  @nonobjc
  public func compileAndRun(
    _ hyloFilePath: String, withOptimizations: Bool, extending p: TypedProgram,
    expecting expectation: ExpectedTestOutcome
  ) throws {
    try checkAnnotatedHyloFileDiagnostics(
      inFileAt: hyloFilePath, expecting: expectation
    ) { (hyloSource, log) in
      try compileAndRun(
        hyloSource, withOptimizations: withOptimizations, extending: p,
        reportingDiagnosticsTo: &log)
    }
  }

  /// Compiles and runs `hyloSource`, applying program optimizations iff `withOptimizations` is
  /// `true`, and `XCTAssert`ing that diagnostics and exit codes match annotated expectations.
  private func compileAndRun(
    _ hyloSource: SourceFile, withOptimizations: Bool, extending baseProgram: TypedProgram,
    reportingDiagnosticsTo log: inout DiagnosticSet
  ) throws {
    var options = ["--emit", "binary"]
    if withOptimizations { options.append("-O") }

    let compilation = try Driver.compileToTemporary(
      hyloSource.url, withOptions: options, extending: baseProgram)
    log.formUnion(compilation.diagnostics)
    try compilation.diagnostics.throwOnError()
    _ = try Process.run(compilation.output, arguments: [])
  }

  /// Compiles the hylo file at `hyloFilePath` up until emitting LLVM code, `XCTAssert`ing that
  /// diagnostics and exit codes match annotated expectations.
  @nonobjc
  public func compileToLLVM(
    _ hyloFilePath: String, extending p: TypedProgram, expecting expectation: ExpectedTestOutcome
  ) throws {
    try checkAnnotatedHyloFileDiagnostics(
      inFileAt: hyloFilePath, expecting: expectation
    ) { (hyloSource, log) in
      let options = ["--emit", "llvm"]
      let compilation = try Driver.compileToTemporary(
        hyloSource.url, withOptions: options, extending: p)
      log.formUnion(compilation.diagnostics)
    }
  }

  /// Lowers the hylo file at `hyloFilePath` to IR, applying any mandatory passes, and `XCTAssert`s
  /// that diagnostics and thrown errors match annotated expectations.
  @nonobjc
  public func lowerToFinishedIR(
    _ hyloFilePath: String, extending p: TypedProgram, expecting expectation: ExpectedTestOutcome
  ) throws {
    try checkAnnotatedHyloFileDiagnostics(
      inFileAt: hyloFilePath, expecting: expectation
    ) { (hyloSource, log) in
      _ = try hyloSource.typecheckedAsMainWithHostedStandardLibrary(reportingDiagnosticsTo: &log, withBuiltinModuleAccess: true)
        .loweredToIR(reportingDiagnosticsTo: &log)
    }
  }

  /// Parses the hylo file at `hyloFilePath`, `XCTAssert`ing that diagnostics and thrown
  /// errors match annotated expectations.
  @nonobjc
  public func parse(
    _ hyloFilePath: String, extending p: TypedProgram, expecting expectation: ExpectedTestOutcome
  ) throws {
    try checkAnnotatedHyloFileDiagnostics(
      inFileAt: hyloFilePath, expecting: expectation
    ) { (hyloSource, log) in
      var ast = AST()
      _ = try ast.loadModule(
        hyloSource.baseName, parsing: [hyloSource], reportingDiagnosticsTo: &log)
    }
  }

  /// Type-checks the Hylo file at `hyloFilePath`, `XCTAssert`ing that diagnostics and thrown
  /// errors match annotated expectations.
  @nonobjc
  public func typeCheck(
    _ hyloFilePath: String, extending p: TypedProgram, expecting expectation: ExpectedTestOutcome
  ) throws {
    try checkAnnotatedHyloFileDiagnostics(
      inFileAt: hyloFilePath, expecting: expectation
    ) { (hyloSource, log) in
      _ = try p.loadModule(reportingDiagnosticsTo: &log) { (ast, log, space) in
        try ast.loadModule(
          hyloSource.baseName, parsing: [hyloSource], inNodeSpace: space,
          withBuiltinModuleAccess: true,
          reportingDiagnosticsTo: &log)
      }
    }
  }

}
